metal: DSpark drafter working end-to-end (3 ops + 4 continuous/drafter fixes)#2
metal: DSpark drafter working end-to-end (3 ops + 4 continuous/drafter fixes)#2robotnursenyc wants to merge 2 commits into
Conversation
…oncat, markov_step) Ports the CUDA kernels dspark_capture_mean_kernel, dspark_gather_concat_kernel, and dspark_markov_step_kernel to Metal, replacing the CUDA-only stubs in ds4_metal.m. Shapes, dtypes, and the markov argmax tie-break (strict-greater scan, lowest index wins) match the CUDA originals; markov_step clamps its threadgroup to the largest power of two <= maxTotalThreadsPerThreadgroup. With these in place the drafter loads, arms, and validates (D1 gate) on a Metal build; the spec loop itself is still gated by the separate cont-prefill nonfinite-logits issue on Metal. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Follow-up: we spent the evening root-causing bug 1 (cont prefill nonfinite seed logits at n >= 8 on Metal) and can narrow it considerably for you, though we have not fully killed it. Findings, all on M3 Ultra, q4 + q2 targets: Where it lives. The cont serving prefill (pos0=0) takes the What we ruled out empirically (each with the n=8/n=30 repro):
Two observations that should localize it for you:
Separate real bug found on the way (independent of the NaN): on the zero_prefix path, Happy to run any diagnostic build you want on the M3 Ultra — this box reproduces in seconds. |
Builds on the three DSpark Metal ops in this branch. Four correctness fixes on the continuous-batching + drafter path: - compressor_update_tensor: the banked emit wrote F32 pool/rms/rope output through an sizeof(float)-strided view of the F16 comp cache, corrupting one element per comp row -> 100% NaN at the first compressed-KV attention layer once n_comp crosses the boundary. Caller-scoped output_is_f16: the banked (F16-cache) caller emits into an F32 scratch then copy_f32_to_f16 into the row; every F32 caller is unchanged. - ds4_gpu_matmul_f32_tensor was GEMV-only (n_tok!=1 returned failure); the DSpark HC-mixer weight is F32 and the block runs it at n_tok=block size. Run n_tok independent GEMVs over per-row views; byte-identical for n_tok==1. - ds4_gpu_tensor_copy blits into the current batch cb and fails when none is open; the DSpark inject block ends the cb, so mtp_cont_rollback_restore_all's restore copies failed. Open a cb around the rollback only when none is open (begin_commands returns 0 when already open -> MTP path unchanged). - HC split/sinkhorn sigmoid: 0.5*tanh(0.5*z)+0.5 overflows to NaN under Metal fast-math tanh (exp-based) for large-positive z. Clamp the input to [-30,30] before the transcendental; sigmoid is saturated by +/-30 so in-range values are unchanged. See METAL_DSPARK.md for details, credits, and observed (content-dependent) acceptance. Root-cause analysis and fixes produced with AI assistance (Claude) and reviewed before submission. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
robotnursenyc
changed the title
2 participants
Makes the DSpark block drafter run end-to-end on Metal for DeepSeek-V4-Flash.
Two parts:
capture_mean,gather_concat,markov_step(original scope of this PR).compressor_update_tensor— the banked emit wrote F32 pool/rms/rope output through asizeof(float)-strided view of the F16 comp cache, corrupting one element per comp row → 100% NaN at the first compressed-KV attention layer oncen_compcrosses the boundary. Fixed with a caller-scopedoutput_is_f16flag (banked/F16 caller emits into an F32 scratch thencopy_f32_to_f16; F32 callers unchanged).ds4_gpu_matmul_f32_tensorwas GEMV-only (n_tok!=1→ failure); the DSpark HC-mixer weight is F32 atn_tok=block size. Now runs per-row GEMVs; byte-identical forn_tok==1.ds4_gpu_tensor_copyneeds an open batch cb; the DSpark inject block ends it, so the rollback's restore copies failed. Open a cb around the rollback only when none is open (MTP path unchanged).0.5*tanh(0.5*z)+0.5overflows to NaN under Metal fast-mathtanhfor large-positivez. Clamp input to[-30,30](saturated region → in-range values unchanged).With all applied, the continuous path serves without serial fallback and the drafter produces real, accepted drafts. Acceptance is content-dependent (~68% / ~2.4× on repetitive text down to ~13% / ~1.2× on free prose, single-run, M3 Ultra, public 2-bit DSpark drafter) — see
METAL_DSPARK.mdfor details, measurements, and credits.Builds clean on Metal. Root-cause analysis and fixes produced with AI assistance (Claude) and reviewed before submission; measurements are single-machine and should be reproduced independently.